/*
* 合成一张复合公钥证书（draft-ietf-lamps-pq-composite-pubkeys）
* ML-DSA-87/ML-KEM-1024
* */

import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyPairGenerator;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPrivateKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPublicKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumSigner;
import org.bouncycastle.pqc.crypto.mlkem.*;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemWriter;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.math.BigInteger;
import java.security.cert.X509Certificate;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Date;
import javax.crypto.*;
import javax.crypto.spec.*;
import java.security.*;


import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.DERSequence;
import org.bouncycastle.pqc.asn1.*;     // 复合公钥辅助类
import org.bouncycastle.asn1.x509.*;
import org.bouncycastle.asn1.x509.KeyUsage;
import org.bouncycastle.asn1.x509.BasicConstraints;
import org.bouncycastle.asn1.x509.ExtendedKeyUsage;
import org.bouncycastle.asn1.x509.KeyPurposeId;

public class Main {
    //注册加密提供者普通密码学和后量子密码学
    static {
        Security.addProvider(new BouncyCastleProvider());
        Security.addProvider(new BouncyCastlePQCProvider());
    }
    // NIST 官方 OID（FIPS 203）
    static ASN1ObjectIdentifier MlKem512 = new ASN1ObjectIdentifier("2.16.840.1.101.3.4.4.1");
    static ASN1ObjectIdentifier MlKem768 = new ASN1ObjectIdentifier("2.16.840.1.101.3.4.4.2");
    static ASN1ObjectIdentifier MlKem1024 = new ASN1ObjectIdentifier("2.16.840.1.101.3.4.4.3");

    public static void main(String[] args) throws Exception {
        try {
            System.out.println("使用Dilithium + CRYSTALS-Kyber 生成复合公钥证书...");
            // 测试ML-DSA-44，ML-DSA-65和ML-DSA-87
            System.out.println("=== ML-DSA-44 + ML-KEM-512 生成X509复合公钥证书示例 ===");
            generateCertificateWithLowLevelAPI(2,MlKem512);

            System.out.println("=== ML-DSA-65 + ML-KEM-768 生成X509复合公钥证书示例 ===");
            generateCertificateWithLowLevelAPI(3,MlKem768);

            System.out.println("=== ML-DSA-87 + ML-KEM-1024 生成X509复合公钥证书示例 ===");
            generateCertificateWithLowLevelAPI(5,MlKem1024);
        } catch (Exception e) {
            System.out.println("证书生成失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private static void generateCertificateWithLowLevelAPI(Integer KeyType,ASN1ObjectIdentifier MLKEYPubkeyType) throws Exception {

        SecureRandom random = new SecureRandom();

        // ML-DSA默认使用Dilithium2参数
        DilithiumParameters params = DilithiumParameters.dilithium2;
        if(KeyType == 2){
            params = DilithiumParameters.dilithium2;
        }
        if(KeyType == 3){
            params = DilithiumParameters.dilithium3;
        }
        if(KeyType == 5){
            params = DilithiumParameters.dilithium5;
        }

        //ML-KEM默认使用ML-KEM-512
        var MLKEYBitSize=MLKEMParameters.ml_kem_512;
        if(MLKEYPubkeyType == MlKem512){
            MLKEYBitSize = MLKEMParameters.ml_kem_512;
        }
        if(MLKEYPubkeyType == MlKem768){
            MLKEYBitSize = MLKEMParameters.ml_kem_768;
        }
        if(MLKEYPubkeyType == MlKem1024){
            MLKEYBitSize = MLKEMParameters.ml_kem_1024;
        }

        // 1) 签名密钥对
        DilithiumKeyPairGenerator keyPairGen = new DilithiumKeyPairGenerator();
        keyPairGen.init(new org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyGenerationParameters(random, params));
        var keyPair = keyPairGen.generateKeyPair();
        DilithiumPublicKeyParameters pubKey = (DilithiumPublicKeyParameters) keyPair.getPublic();
        DilithiumPrivateKeyParameters privKey = (DilithiumPrivateKeyParameters) keyPair.getPrivate();
        System.out.println("Dilithium"+ KeyType +" 密钥对生成成功");

        // 2) 封装密钥对（ML-KEM-512,ML-KEM-768,ML-KEM-1024） 根据MLKEYBitSize指定
        MLKEMKeyPairGenerator kkpGen = new MLKEMKeyPairGenerator();
        kkpGen.init(new MLKEMKeyGenerationParameters(random, MLKEYBitSize));
        AsymmetricCipherKeyPair kKeyPair = kkpGen.generateKeyPair();
        MLKEMPublicKeyParameters kPub  = (MLKEMPublicKeyParameters)  kKeyPair.getPublic();
        MLKEMPrivateKeyParameters kKey  = (MLKEMPrivateKeyParameters)  kKeyPair.getPrivate();

        // 单个公钥的 AlgorithmIdentifier
        AlgorithmIdentifier mlDsaAlgId = new AlgorithmIdentifier(
                getAlgorithmOID(params));
        AlgorithmIdentifier mlKemAlgId = new AlgorithmIdentifier(
                MLKEYPubkeyType);

        PrivateKeyInfo MLDSAprikeypkcs8 = new PrivateKeyInfo(mlDsaAlgId,privKey.getEncoded());
        SubjectPublicKeyInfo MLDSApubkeypkcs8 = new SubjectPublicKeyInfo(mlDsaAlgId,pubKey.getEncoded());

        PrivateKeyInfo MLKEMprikeypkcs8 = new PrivateKeyInfo(mlKemAlgId,kKey.getEncoded());
        SubjectPublicKeyInfo MLKEMpubkeypkcs8 = new SubjectPublicKeyInfo(mlKemAlgId,kPub.getEncoded());

        // 分别编码成 SubjectPublicKeyInfo
        SubjectPublicKeyInfo dSpi = new SubjectPublicKeyInfo(mlDsaAlgId, pubKey.getEncoded());
        SubjectPublicKeyInfo kSpi = new SubjectPublicKeyInfo(mlKemAlgId, kPub.getEncoded());
        // 复合公钥结构：SEQUENCE { SPKI1, SPKI2 }
        ASN1Sequence compositeKeySeq = new DERSequence(new ASN1Encodable[]{dSpi, kSpi});

        // 顶层复合 AlgorithmIdentifier（固定 OID） 根据PQC官方公告，1.3.6.1.4.1.44363.45为复合公钥实验性OID
        AlgorithmIdentifier compositeAlgId = new AlgorithmIdentifier(
                new ASN1ObjectIdentifier("1.3.6.1.4.1.44363.45"));       //复合公钥结构  ML-DSA + ML-KEM 的算法OID 1.3.6.1.4.1.44363.45

        // 最终 SubjectPublicKeyInfo  公钥信息 - 使用双公钥复合证书的公钥信息
        SubjectPublicKeyInfo compositePubInfo =
                new SubjectPublicKeyInfo(compositeAlgId, compositeKeySeq.getEncoded());


        // 设置正确的签名算法OID和参数（包含NULL）
        AlgorithmIdentifier sigAlgId = getAlgorithmIdentifier(params);
        System.out.println("使用算法OID: " + sigAlgId.getAlgorithm());

        // 构建证书
        X500Name subject = new X500Name("CN=127.0.0.1, OU=Test Dilithium"+ KeyType + " " +MLKEYBitSize.getName() + " Composite Certificate, O=Demo, C=CN");
        BigInteger serial = BigInteger.valueOf(System.currentTimeMillis());
        Instant now = Instant.now();
        Date notBefore = Date.from(now);
        Date notAfter = Date.from(now.plus(365, ChronoUnit.DAYS));




        X509v3CertificateBuilder certGen = new X509v3CertificateBuilder(
                subject, serial, notBefore, notAfter, subject, compositePubInfo);
        /* ************************
        digitalSignature	0	数字签名	TLS/SSL 握手、证书签名、代码签名
        nonRepudiation	    1	不可否认（也称 contentCommitment）	电子签名、时间戳
        keyEncipherment	    2	密钥加密	RSA/ECC 密钥交换（传统 TLS）
        dataEncipherment	3	数据加密	直接用公钥加密数据（罕见）
        keyAgreement	    4	密钥协商	ECDH、ML-KEM 密钥交换
        keyCertSign	        5	证书签名	CA 证书签发下级证书
        cRLSign	            6	CRL 签名	CA 签发证书吊销列表
        encipherOnly	    7	仅加密*	与 keyAgreement 联合使用
        decipherOnly	    8	仅解密*	与 keyAgreement 联合使用
        * 注意：第 7、8 位不能单独出现，必须与 keyAgreement 一起使用，用来限定“只能加密”或“只能解密”。
         *********************** */
        // 添加 密钥用法 KeyUsage 扩展
        KeyUsage keyUsage = new KeyUsage(
                KeyUsage.digitalSignature | KeyUsage.dataEncipherment | KeyUsage.keyAgreement |
                        KeyUsage.keyEncipherment
        );
        certGen.addExtension(
                org.bouncycastle.asn1.x509.Extension.keyUsage,
                true, // critical
                keyUsage
        );

        // 添加 增强型密钥用法 ExtendedKeyUsage 扩展
        ExtendedKeyUsage eku = new ExtendedKeyUsage(new KeyPurposeId[]{
                KeyPurposeId.id_kp_clientAuth,
                KeyPurposeId.id_kp_serverAuth
        });
        certGen.addExtension(
                org.bouncycastle.asn1.x509.Extension.extendedKeyUsage,
                false, // non-critical
                eku
        );
        // 添加 基本约束：CA=false，最终实体证书
        certGen.addExtension(org.bouncycastle.asn1.x509.Extension.basicConstraints, true, new BasicConstraints(false));


        // 自定义内容签名器
        ContentSigner signer = new ContentSigner() {
            private final java.io.ByteArrayOutputStream buffer = new java.io.ByteArrayOutputStream();

            @Override
            public AlgorithmIdentifier getAlgorithmIdentifier() {
                return sigAlgId;
            }

            @Override
            public java.io.OutputStream getOutputStream() {
                return buffer;
            }

            @Override
            public byte[] getSignature() {
                try {
                    byte[] dataToSign = buffer.toByteArray();
                    DilithiumSigner dilithiumSigner = new DilithiumSigner();
                    dilithiumSigner.init(true, privKey);
                    byte[] signature = dilithiumSigner.generateSignature(dataToSign);

                    // Dilithium签名通常已经是ASN.1编码格式，但确保正确
                    return signature;
                } catch (Exception e) {
                    throw new RuntimeException("签名失败: " + e.getMessage(), e);
                }
            }
        };

        System.out.println("开始构建证书...");
        org.bouncycastle.cert.X509CertificateHolder certHolder = certGen.build(signer);
        System.out.println("证书构建完成");

        // 转换为JCA证书
        X509Certificate cert = new JcaX509CertificateConverter()
                .setProvider("BC")
                .getCertificate(certHolder);

        System.out.println("证书生成成功");

        // 验证签名算法标识符包含NULL参数
        byte[] encodedCert = cert.getEncoded();
        // 可以在这里添加验证逻辑来检查是否包含05 00

        // 输出X509复合证书文件
        writeCertificateToPEM(cert, "./out/dilithium"+ KeyType +"CRYSTALSKyber_cert.crt");
        writeCertificateToDER(cert, "./out/dilithium"+ KeyType +"CRYSTALSKyber_cert.der");
        // 输出X509复合证书 签名算法的公钥和私钥
        writePrivateKeyToPEM(MLDSAprikeypkcs8.getEncoded(), "./out/dilithium"+ KeyType +"CRYSTALSKyber_SignPrivate_key.pem");
        writePublicKeyToPEM(MLDSApubkeypkcs8.getEncoded(), "./out/dilithium"+ KeyType +"CRYSTALSKyber_SignPublic_key.pem");
        // 输出X509复合证书 密钥封装算法的公钥和私钥
        writePrivateKeyToPEM(MLKEMprikeypkcs8.getEncoded(), "./out/dilithium"+ KeyType +"CRYSTALSKyber_KeyEncapsulationPrivate_key.pem");
        writePublicKeyToPEM(MLKEMpubkeypkcs8.getEncoded(), "./out/dilithium"+ KeyType +"CRYSTALSKyber_KeyEncapsulationPublic_key.pem");

        System.out.println("=== 证书信息 ===");
        System.out.println("主题: " + cert.getSubjectX500Principal());
        System.out.println("颁发者: " + cert.getIssuerX500Principal());
        System.out.println("签名算法: " + cert.getSigAlgName());
        System.out.println("序列号: " + cert.getSerialNumber());
        System.out.println("有效期: " + cert.getNotBefore() + " 至 " + cert.getNotAfter());
        System.out.println("所有文件已成功生成！");
    }

    private static AlgorithmIdentifier getAlgorithmIdentifier(DilithiumParameters params) {
        // 返回包含NULL参数的算法标识符
        ASN1ObjectIdentifier oid = getAlgorithmOID(params);
        return new AlgorithmIdentifier(oid); // 添加NULL参数, DERNull.INSTANCE
    }

    private static ASN1ObjectIdentifier getAlgorithmOID(DilithiumParameters params) {
        // 根据Dilithium参数返回正确的OID
        if (params == DilithiumParameters.dilithium2) {
            return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.17");
        } else if (params == DilithiumParameters.dilithium3) {
            return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.18");
        } else if (params == DilithiumParameters.dilithium5) {
            return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.19");
        } else {
            return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.19");
        }
    }


    // PEM输出方法保持不变
    private static void writeCertificateToPEM(X509Certificate cert, String filename) throws Exception {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("CERTIFICATE", cert.getEncoded());
            pemWriter.writeObject(pemObject);
        }
        System.out.println("证书已保存: " + filename);
    }

    private static void writePrivateKeyToPEM(byte[] privateKeyBytes, String filename) throws IOException {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("PRIVATE KEY", privateKeyBytes);
            pemWriter.writeObject(pemObject);
        }
        System.out.println("私钥已保存: " + filename);
    }

    private static void writePublicKeyToPEM(byte[] publicKeyBytes, String filename) throws IOException {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("PUBLIC KEY", publicKeyBytes);
            pemWriter.writeObject(pemObject);
        }
        System.out.println("公钥已保存: " + filename);
    }

    private static void writeCertificateToDER(X509Certificate cert, String filename) throws Exception {
        try (FileOutputStream fos = new FileOutputStream(filename)) {
            fos.write(cert.getEncoded());
        }
        System.out.println("DER证书已保存: " + filename);
    }
}
